#include "task_wrapper.hpp"
#include <iostream>
#include <chrono>
#include <thread>
#include <cassert>
#include <string>

using namespace timeout_manager;

// g++ -std=c++17 -pthread -o test_task_wrapper test_task_wrapper.cpp

// 测试函数：简单计算
int simpleCalculation(int a, int b) {
    std::cout << "执行计算: " << a << " + " << b << std::endl;
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    return a + b;
}

// 测试函数：无返回值
void printMessage(const std::string& msg) {
    std::cout << "消息: " << msg << std::endl;
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
}

// 测试函数：会抛出异常
int throwingFunction() {
    std::cout << "准备抛出异常..." << std::endl;
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    throw std::runtime_error("测试异常");
    return 42;  // 永远不会执行到这里
}

// 测试函数：长时间运行，支持取消检查
int longRunningTask(int duration, const TaskWrapper<int>* wrapper) {
    std::cout << "开始长时间任务，持续 " << duration << "ms" << std::endl;
    
    int progress = 0;
    const int stepDuration = 50;  // 每50ms检查一次
    
    while (progress < duration) {
        // 检查是否应该取消
        if (wrapper && wrapper->shouldCancel()) {
            std::cout << "任务检测到取消信号，提前退出" << std::endl;
            throw TimeoutException("任务被用户取消");
        }
        
        std::this_thread::sleep_for(std::chrono::milliseconds(stepDuration));
        progress += stepDuration;
        
        if (progress % 200 == 0) {  // 每200ms输出进度
            std::cout << "任务进度: " << progress << "/" << duration << "ms" << std::endl;
        }
    }
    
    std::cout << "长时间任务完成" << std::endl;
    return duration;
}

void testBasicExecution() {
    std::cout << "\n=== 测试基本任务执行 ===" << std::endl;
    
    // 测试有返回值的任务
    {
        // 1. 获取函数包装器
        auto wrapper = std::make_unique<TaskWrapper<int>>(1, 
            []() { return simpleCalculation(10, 20); });
        
        // 2. 获取函数运行结果的future对象
        auto future = wrapper->getFuture();
        
        // 检查初始状态
        assert(wrapper->getStatus() == TaskStatus::PENDING);
        std::cout << "✓ 初始状态正确" << std::endl;
        
        // 3. 在新线程中调用execute运行原函数
        std::thread executor([&wrapper]() {
            wrapper->execute();
        });
        
        // 4. 获取结果
        int result = future.get();
        std::cout << "任务结果: " << result << std::endl;
        assert(result == 30);
        
        // 5. 等待线程结束
        executor.join();
        assert(wrapper->getStatus() == TaskStatus::COMPLETED);
        std::cout << "✓ 有返回值任务执行正确" << std::endl;
    }
    
    // 测试无返回值的任务
    {
        // 1. 获取函数包装器
        auto wrapper = std::make_unique<TaskWrapper<void>>(2, 
            []() { printMessage("Hello from wrapper!"); });
        
        // 2. 获取函数运行结果的future对象
        auto future = wrapper->getFuture();
        
        // 3. 在新线程中调用execute运行原函数
        std::thread executor([&wrapper]() {
            wrapper->execute();
        });
        
        // 4. 等待完成
        future.get();  
        
        // 5. 等待线程结束
        executor.join(); 
        assert(wrapper->getStatus() == TaskStatus::COMPLETED);
        std::cout << "✓ 无返回值任务执行正确" << std::endl;
    }
}

void testTaskCancellation() {
    std::cout << "\n=== 测试任务取消 ===" << std::endl;
    
    // 测试执行前取消
    {
        auto wrapper = std::make_unique<TaskWrapper<int>>(10, 
            []() { return simpleCalculation(5, 10); });
        
        auto future = wrapper->getFuture();
        
        // 立即取消任务
        wrapper->cancel();
        assert(wrapper->getStatus() == TaskStatus::CANCELLED);
        std::cout << "✓ 执行前取消状态正确" << std::endl;
        
        // 尝试执行已取消的任务
        wrapper->execute();
        
        try {
            future.get();
            assert(false);  // 不应该到达这里
        } catch (const TimeoutException& e) {
            std::cout << "✓ 正确捕获取消异常: " << e.what() << std::endl;
        }
    }
    
    // 测试执行过程中取消（需要任务配合检查shouldCancel）
    {
         // 首先定义wrapper为nullptr
        std::shared_ptr<TaskWrapper<int>> wrapper;
        
        // 定义lambda函数，通过引用捕获wrapper
        auto task = [&wrapper]() -> int {
            return longRunningTask(500, wrapper.get());
        };
    
        // 然后创建包装器
        wrapper = std::make_shared<TaskWrapper<int>>(11, std::function<int()>(task));
        
        auto future = wrapper->getFuture();
        
        // 在新线程中执行任务
        std::thread executor([wrapper]() {
            wrapper->execute();
        });
        
        // 等待一小段时间后取消
        std::this_thread::sleep_for(std::chrono::milliseconds(250));
        wrapper->cancel();
        
        try {
            future.get();
            assert(false);  // 不应该到达这里
        } catch (const TimeoutException& e) {
            std::cout << "✓ 正确捕获执行中取消异常: " << e.what() << std::endl;
        }
        
        executor.join();
    }
}

void testTaskTimeout() {
    std::cout << "\n=== 测试任务超时 ===" << std::endl;
    
    auto wrapper = std::make_unique<TaskWrapper<int>>(20, 
        []() { 
            std::this_thread::sleep_for(std::chrono::milliseconds(2000));
            return 100; 
        });
    
    auto future = wrapper->getFuture();
    
    // 在新线程中执行任务
    std::thread executor([&wrapper]() {
        wrapper->execute();
    });
    
    // 等待一小段时间后设置超时
    std::this_thread::sleep_for(std::chrono::milliseconds(1000));
    wrapper->timeout();
    
    assert(wrapper->getStatus() == TaskStatus::TIMEOUT);
    std::cout << "✓ 超时状态设置正确" << std::endl;
    
    try {
        future.get();
        assert(false);  // 不应该到达这里
    } catch (const TimeoutException& e) {
        std::cout << "✓ 正确捕获超时异常: " << e.what() << std::endl;
    }
    
    executor.join();
}

void testExceptionHandling() {
    std::cout << "\n=== 测试异常处理 ===" << std::endl;
    
    auto wrapper = std::make_unique<TaskWrapper<int>>(30, throwingFunction);
    auto future = wrapper->getFuture();
    
    std::thread executor([&wrapper]() {
        wrapper->execute();
    });
    
    try {
        future.get();
        assert(false);  // 不应该到达这里
    } catch (const std::runtime_error& e) {
        std::cout << "✓ 正确捕获任务异常: " << e.what() << std::endl;
    }
    
    executor.join();
    assert(wrapper->getStatus() == TaskStatus::COMPLETED);
    std::cout << "✓ 异常任务状态正确" << std::endl;
}

void testStatusTransitions() {
    std::cout << "\n=== 测试状态转换 ===" << std::endl;
    
    auto wrapper = std::make_unique<TaskWrapper<int>>(40, 
        []() { 
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            return 42; 
        });
    
    // 初始状态
    assert(wrapper->getStatus() == TaskStatus::PENDING);
    std::cout << "✓ 初始状态: PENDING" << std::endl;
    
    auto future = wrapper->getFuture();
    
    std::thread executor([&wrapper]() {
        wrapper->execute();
    });
    
    // 等待任务开始执行
    std::this_thread::sleep_for(std::chrono::milliseconds(10));
    
    // 此时应该是RUNNING状态
    auto status = wrapper->getStatus();
    std::cout << "执行中状态: " << taskStatusToString(status) << std::endl;
    
    // 等待任务完成
    int result = future.get();
    executor.join();
    
    assert(wrapper->getStatus() == TaskStatus::COMPLETED);
    assert(result == 42);
    std::cout << "✓ 最终状态: COMPLETED" << std::endl;
}

int main() {
    std::cout << "开始测试任务包装器..." << std::endl;
    
    testBasicExecution();
    testTaskCancellation();
    testTaskTimeout();
    testExceptionHandling();
    testStatusTransitions();
    
    std::cout << "\n所有任务包装器测试通过！" << std::endl;
    return 0;
}